from python import Python
from dose_stat import DoseStatRet, accumulate, get_metric, get_dose_stat_roi_interp

# 将Mojo类型转换为Python对象
fn dose_stat_ret_to_python(result: DoseStatRet) -> PythonObject:
    var py_dict = Python.dict()
    py_dict["binned_dose_values"] = Python.list(result.binned_dose_values)
    py_dict["accumulated_relative_volumes"] = Python.list(result.accumulated_relative_volumes)
    return py_dict

# Python兼容的接口函数
@value
struct DoseStatisticsPy:
    fn accumulate_py(
        voxel_volume: PythonObject,
        voxel_dose: PythonObject,
        do_exact_accumulation: Bool = False
    ) -> PythonObject:
        # 转换Python列表到Mojo列表
        var volume_list = List[Float64]()
        var dose_list = List[Float64]()
        
        for i in range(len(voxel_volume)):
            volume_list.append(Float64(voxel_volume[i]))
            dose_list.append(Float64(voxel_dose[i]))
        
        # 调用Mojo函数
        var result = accumulate(volume_list, dose_list, do_exact_accumulation)
        return dose_stat_ret_to_python(result)
    
    fn get_metric_py(
        dose_at_volume: Float64,
        binned_dose_values: PythonObject,
        accumulated_relative_volumes: PythonObject
    ) -> Float64:
        var binned_list = List[Float64]()
        var volumes_list = List[Float64]()
        
        for i in range(len(binned_dose_values)):
            binned_list.append(Float64(binned_dose_values[i]))
        
        for i in range(len(accumulated_relative_volumes)):
            volumes_list.append(Float64(accumulated_relative_volumes[i]))
        
        return get_metric(dose_at_volume, binned_list, volumes_list)
    
    fn get_dose_stat_roi_interp_py(
        dose_type: String,
        binned_dose_values: PythonObject,
        accumulated_relative_volumes: PythonObject,
        voxel_volume: PythonObject,
        voxel_dose: PythonObject,
        dose_at_volume: Float64 = 0.0
    ) -> Float64:
        var binned_list = List[Float64]()
        var volumes_list = List[Float64]()
        var volume_list = List[Float64]()
        var dose_list = List[Float64]()
        
        # 转换所有Python列表到Mojo列表
        for i in range(len(binned_dose_values)):
            binned_list.append(Float64(binned_dose_values[i]))
        
        for i in range(len(accumulated_relative_volumes)):
            volumes_list.append(Float64(accumulated_relative_volumes[i]))
        
        for i in range(len(voxel_volume)):
            volume_list.append(Float64(voxel_volume[i]))
            dose_list.append(Float64(voxel_dose[i]))
        
        return get_dose_stat_roi_interp(
            dose_type, binned_list, volumes_list, volume_list, dose_list, dose_at_volume
        )

# 导出到Python模块
fn initialize_module():
    var dose_stats = DoseStatisticsPy()
    
    # 创建模块字典
    var module_dict = Python.dict()
    module_dict["accumulate"] = dose_stats.accumulate_py
    module_dict["get_metric"] = dose_stats.get_metric_py
    module_dict["get_dose_stat_roi_interp"] = dose_stats.get_dose_stat_roi_interp_py
    
    return module_dict

# 模块入口点
fn main():
    var module = initialize_module()
    print("剂量统计Python模块已初始化")